#include "TestCase.h"
#include "JSON.h"
#include <string>
#include <fstream>
#include <sstream>
#include <cctype>
#include <cstring>
#include <iostream>
#include <unordered_map>
#include <algorithm>
using namespace std;

namespace {
  /* Type representing the data from the student payload. */
  struct StudentData {
    vector<string> states; // States visited, in order
    string emailAddress;   // Email address
    size_t hash;           // Hash code for the data payload
  };

  /* Generic "something is wrong with the payload" message. */
  const string kMalformedCode =
    "The code that you submitted appears to be malformed. Please make sure that you "
    "copied all of the gold text and didn't copy any of the other text with it.";
    
  /* Given a string, converts it to an integer, failing if it's malformed. */
  size_t stringToInteger(const string& str) {
    istringstream extractor(str);
    
    size_t result;
    if (extractor >> noskipws >> result, !extractor) failTestVisibly(kMalformedCode);
    
    char leftover;
    if (extractor >> leftover) failTestVisibly(kMalformedCode);
    
    return result;
  }
  
  const int HASH_SEED = 5381;               // Starting point for first cycle
  const int HASH_MULTIPLIER = 33;           // Multiplier for each cycle
  const int HASH_MASK = unsigned(-1) >> 1;  // All 1 bits except the sign

  size_t hashCode(int key) {
      return key & HASH_MASK;
  }

  size_t hashCode(const char* base, size_t numBytes) {
    unsigned hash = HASH_SEED;
    for (size_t i = 0; i < numBytes; i++) {
      hash = HASH_MULTIPLIER * hash + base[i];
    }
    return hashCode(hash);
  }
  
  size_t hashCode(const string& str) {
    return hashCode(str.data(), str.length());
  }
  
  /* Constructs the map used to decode base64-encoded text. */
  unordered_map<char, unsigned> decodingMap() {
    const string kBase64Chars = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
    
    unordered_map<char, unsigned> result;
    for (size_t i = 0; i < kBase64Chars.size(); i++) {
      result[kBase64Chars[i]] = i;
    }
    if (result.size() != 64) internalError("Wrong number of characters in Base64 map?");
    
    return result;
  }
  
  /* Given four base64 characters, decodes them to three logical characters. */
  string decodeBlock(const string& block,
                     const unordered_map<char, unsigned>& table) {
    if (block.size() != 4) internalError("Block must have length four.");
    
    /* Place all the bits in the right places. */
    uint32_t bits = 0;
    for (size_t i = 0; i < 4; i++) {
      if (!table.count(block[i])) {
        cout << "ERROR: Unknown character: 0x" << hex << static_cast<unsigned int>(static_cast<unsigned char>(block[i])) << endl;
        failTestVisibly(kMalformedCode);
      }
      
      bits += table.at(block[i]) << (24 - 6 * (i + 1));
    }
    
    /* Convert the bits to characters. */
    string result;
    for (size_t i = 0; i < 3; i++) {
      result += char(bits >> ((3 - (i + 1)) * 8));
    }
    return result;
  }
  
  /* Decodes a piece of text that's Base64-encoded. */
  string base64Decode(const string& input) {
    /* Build a map from characters to bit patterns. */
    auto charToValue = decodingMap();
    
    /* Confirm the length is indeed a multiple of 4. */
    if (input.size() % 4 != 0) {  
      cout << "ERROR: Data length isn't a multiple of four. " << endl;
      failTestVisibly(kMalformedCode);
    }
    
    /* Count how many trailing equals signs there are, replacing each one with an A
     * (decodes to 0). This allows us to decode everything all at once.
     */
    string toDecode = input;
    size_t padding = 0;
    
    for (size_t i = toDecode.size(); i > 0; i--) {
      if (toDecode[i - 1] == '=') {
        toDecode[i - 1] = 'A';
        padding++;
      }
    }
    if (padding > 2) {
      cout << "ERROR: Too many padding bytes." << endl;
      failTestVisibly(kMalformedCode);
    }
    
    /* Decode everything one block at a time. */
    string result;
    for (size_t i = 0; i < toDecode.size(); i += 4) {
      result += decodeBlock(string(toDecode.begin() + i, toDecode.begin() + i + 4), charToValue);
    }
    
    /* Trim off the padding characters and return everything else. */
    return result.substr(0, result.size() - padding);
  }
  

  /* Loads the student's email address from the test metadata. */
  string studentEmailAddress() {
    ifstream input("/autograder/submission_metadata.json");
    if (!input) internalError("Cannot open metadata file.");
  
    JSON users = JSON::parse(input)["users"];
    if (users.size() != 1) internalError("Should be exactly one submitter?");
    
    return users[0]["email"].asString();
  }
  
  /* Returns the contents of the student's submitted code. */
  string studentCode() {
    ifstream input("code.txt");
    if (!input) internalError("Cannot open code.txt for reading.");
    
    string result;
    for (string line; getline(input, line); ) {
      /* Copy all non-whitespace characters over. */
      for (char ch: line) {
        if (!isspace(ch)) result += ch;
      }
    }
    
    return result;
  }
  
  /* Returns the raw string data in the payload from the submitted code. */
  const size_t kHeaderSize = 16;
  string rawStudentData() {
    /* Begin by base64-decoding the student's code. */
    string data = base64Decode(studentCode());
    
    /* This needs to be long enough for there to be a header. */
    if (data.size() < kHeaderSize) {
      cout << "ERROR: Data is too short." << endl;
      failTestVisibly(kMalformedCode);
    }
    
    /* Decode the remainder of the payload. */
    string payload;
    for (size_t i = kHeaderSize; i < data.size(); i++) {
      payload += char(data[i] ^ data[i % kHeaderSize]);
    }
    
    cout << "Payload is " << payload << endl;
    return payload;
  }
  
  /* Returns the payload from the student's submitted code. */
  StudentData studentPayload() {
    StudentData result;
    
    /* Split the data at slash boundaries. */
    string data = rawStudentData();
    
    size_t i = 0;
    while (true) {
      size_t next = data.find('/', i);
      if (next == string::npos) {
        result.states.push_back(data.substr(i));
        break;
      }
      
      /* Grab everything from the current point to the slash. */
      result.states.push_back(string(data.begin() + i, data.begin() + next));
      
      /* Skip over the slash. */
      i = next + 1;
    }
    
    /* The last two items should be the email and the validation code. Move those
     * over.
     */
    if (result.states.size() < 2) {
      cout << "ERROR: Too few states." << endl;
      failTestVisibly(kMalformedCode);
    }
    
    result.emailAddress = result.states[result.states.size() - 2];
    result.hash = stringToInteger(result.states[result.states.size() - 1]);
    
    result.states.erase(result.states.begin() + result.states.size() - 2, result.states.end());
    return result;
  }
  
  /* Given the student's data, which represents the path they took, make sure they hit the points
   * that we needed them to hit.
   */
  void confirmValidPath(const vector<string>& data) {
    /* Make sure we start and end in the right place. */
    if (data.front() != "Welcome" ||
        data.back()  != "ThatsAWrap") {
        cout << "ERROR: Start and/or end point is wrong." << endl;
        failTestVisibly(kMalformedCode);    
    }
    
    /* Make sure we see each Honor Code point hit. */
    for (size_t i = 1; i <= 10; i++) {
      if (find(data.begin(), data.end(), "HonorCode" + to_string(i) + "Done") == data.end()) {
        cout << "ERROR: Missing Honor Code question " << i << endl;
        failTestVisibly(kMalformedCode);
      }
    }
  }
  
  /* Checks whether an email address is structurally valid. This uses the automaton
   * described in lecture!
   */
  bool isValidEmailAddress(const string& email) {
    int state = 0;
    
    for (char ch: email) {
      if (isalnum(ch)) {
        if      (state == 0) state = 1;
        else if (state == 1) state = 1;
        else if (state == 2) state = 1;
        else if (state == 3) state = 4;
        else if (state == 4) state = 4;
        else if (state == 5) state = 6;
        else if (state == 6) state = 6;
        else if (state == 7) state = 6;
        else internalError("What state are we in?");
      } else if (ch == '.') {
        if      (state == 0) return false;
        else if (state == 1) state = 2;
        else if (state == 2) return false;
        else if (state == 3) return false;
        else if (state == 4) state = 5;
        else if (state == 5) return false;
        else if (state == 6) state = 7;
        else if (state == 7) return false;
        else internalError("What state are we in?");
      } else if (ch == '@') {
        if      (state == 0) return false;
        else if (state == 1) state = 3;
        else if (state == 2) return false;
        else if (state == 3) return false;
        else if (state == 4) return false;
        else if (state == 5) return false;
        else if (state == 6) return false;
        else if (state == 7) return false;
        else internalError("What state are we in?");
      } else {
        return false;
      }
    }
    
    return state == 6;
  }
  
  /* Confirms that the student payload email address matches the account email address. */
  void confirmEmailAddress(const StudentData& data) {
    string expected = studentEmailAddress();
    
    /* Make sure that this account ends with "@stanford.edu" */
    size_t at = expected.find('@');
    if (at == string::npos) internalError("No at sign in email address?");
    
    string domain = expected.substr(at + 1);
    if (domain != "stanford.edu") {
      ostringstream builder;
      builder << "The email address you used to register for this GradeScope account ("
              << expected
              << ") isn't an @stanford.edu account. For internal record-keeping purposes, "
              << "we need each GradeScope account to be set up with an @stanford.edu account. "
              << "Please delete this account and create a new one using your @stanford.edu "
              << "email address. Sorry for the inconvience!";
      failTestVisibly(builder.str());              
    }
    
    if (!isValidEmailAddress(data.emailAddress)) {
      cout << "ERROR: Email address is malformed." << endl;
      failTestVisibly(kMalformedCode);
    }
    
    if (data.emailAddress != expected) {
      ostringstream builder;
      builder << "The email address associated with this code ("
              << data.emailAddress
              << ") doesn't match the email address associated with this GradeScope account ("
              << expected
              << "). Make sure that you entered your email address correctly at the end of "
              << "the assignment.";
      failTestVisibly(builder.str());
    }
  }
  
  void confirmValidHash(const StudentData& data) {
    /* Rebuild the slash-separated string. */
    string rawBytes;
    
    for (const string& state: data.states) {
      rawBytes += state + "/";
    }
    
    rawBytes += data.emailAddress + "/";
    
    if (hashCode(rawBytes) != data.hash) {
      cout << "ERROR: Hash is incorrect." << endl;
      failTestVisibly(kMalformedCode);
    }
  }
}

ADD_TEST("Code Verification", 5) {
  auto data = studentPayload();
  
  /* Confirm that the hash is valid. */
  confirmValidHash(data);
  
  /* The list of states should have at least two items in it. */
  if (data.states.size() < 2) {
    cout << "ERROR: Too few states in path." << endl;
    failTestVisibly(kMalformedCode);
  }

  /* Make sure that the path encoded here is a valid one. */
  confirmValidPath(data.states);
  
  /* Make sure that the email address matches what we expect. */
  confirmEmailAddress(data);
}
